import PyPDF2
from reportlab.lib.pagesizes import letter
from reportlab.lib.units import inch
from reportlab.pdfgen import canvas
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.pdfbase import pdfmetrics
from datetime import datetime
import os 

from typing import List
from pikepdf import Pdf,Page,Rectangle


'''
向目标pdf文件批量添加水印
target_pdf_path:目标pdf文件路径+文件名
watermark_pad_path:水印pdf文件路径+文件名
nrow:水印平铺的行数
ncol:水印平铺的列数
skip_pages:需要跳过不添加水印的页数

'''
class Watermark:
    def __init__(self, 
                skip_pages:List[int] = [],
            ) -> None:        
        """
        为PDF文件添加水印。

        参数:
        watermark_pdf_path (str): 水印PDF文件的路径。
        nrow (int): 每行添加的水印数量。
        ncol (int): 每列添加的水印数量。
        skip_pages (List[int], 可选): 需要跳过不添加水印的页面编号列表，默认为空列表。

        返回:
        None
        """
        self.watermark_pdf_path = '/root/k8s-vue/exam-platform-of-big-data/工具函数/watermark.pdf'
        self.watermark_score_pdf_path = '/root/k8s-vue/exam-platform-of-big-data/工具函数/watermark_score.pdf'
        self.nrow = 1
        self.ncol = 1
        self.skip_pages = skip_pages

    def  generator_score_watermark(self, teacher, score,max_score, center_x=100, center_y=100) -> None:
        # Register Chinese font
        pdfmetrics.registerFont(TTFont('song', 'simsun.ttc'))

        c = canvas.Canvas(self.watermark_pdf_path, pagesize=(200,200))
        def blue_color():
            c.setStrokeColorRGB(0.1, 0.6, 0.8)
            c.setFillColorRGB(0.1, 0.6, 0.8)
        def red_color():
            c.setStrokeColorRGB(1, 0, 0)
            c.setFillColorRGB(1, 0, 0)
        # Draw watermark circle
        
        red_color()
        c.setLineWidth(width=4)
        c.circle(x_cen=center_x, y_cen=center_y, r=60)
        
        c.setLineWidth(width=1)
        c.circle(x_cen=center_x, y_cen=center_y, r=55)

        # Add watermark text in Chinese
        c.setFont("song", 10)  # Use the registered Chinese font
        c.drawString(center_x-35, center_y+10, '成绩：')
        c.drawString(center_x-35, center_y-10, '阅卷人：')
        
        # Add timestamp
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        c.setFont("song", 9)
        c.drawString(center_x-40, center_y-30, f"{timestamp}")

        # 得分、教师名、★， 用红色表示
        red_color()
        c.drawString(center_x-5, center_y+10, '{}/{}'.format(str(score), str(max_score)))
        c.drawString(center_x, center_y-10, teacher)
        percent = 1.0*score/max_score
        if percent >= 0.8:
            c.drawString(center_x-20, center_y+30, '★★★★★')
        elif percent >= 0.6:
            c.drawString(center_x-20, center_y+30, '★★★★')
        elif percent >= 0.4:
            c.drawString(center_x-20, center_y+30, '★ ★ ★')
        elif percent >= 0.2:
            c.drawString(center_x-20, center_y+30, '  ★ ★')
        else:
            c.drawString(center_x-20, center_y+30, '   ★')


        c.save()
        
    
    def  generator_score_postions_watermark(self,  postions, page, page_width, page_height) -> None:
        # Register Chinese font
        pdfmetrics.registerFont(TTFont('song', 'simsun.ttc'))

        c = canvas.Canvas(self.watermark_score_pdf_path, pagesize=(page_width,page_height))


        for pos in postions:
            if pos['page'] == page:
                #计算水印在目标页面上的位置并添加
                # 颜色改成红色
                c.setFillColorRGB(1, 0, 0)
                # 字体加粗
                c.setFont("song", 10, 2)  # Use the registered Chinese font
        
                x1,y1,x2,y2 = pos['pos'].strip().split(' ')
                c.setFont("song", 10)  # Use the registered Chinese font
                score = pos['score']
                max_score = pos['max_score']
                c.drawString(float(x1),float(y2)-10, f'{score}/{max_score}')

        c.save()
        
    def _add_score_watermark(self, pdf_path ) -> None:

        
        #打开目标PDF文件
        target_pdf = Pdf.open(pdf_path, allow_overwriting_input=True)

        #打开水印PDF文件并提取第一页作为水印
        watermark_pdf = Pdf.open(self.watermark_pdf_path)
        watermark_page = watermark_pdf.pages[0]

        #遍历目标PDF的每一页，根据行列数批量添加水印
        for idx,target_page in enumerate(target_pdf.pages):
            if idx in self.skip_pages:  # 如果当前页码在skip_pages列表中，则跳过
                continue
            for x in range(self.ncol):
                for y in range(self.nrow):
                    #计算水印在目标页面上的位置并添加
                    
                    # target_page.mediabox.lower_left = (0,0)
                    # target_page.mediabox.lowerLeft  = (0,0)
                    target_page.add_overlay(
                        watermark_page,
                        Rectangle(
                            # 右上角水印左下角X坐标，通常设为距离页面右侧边缘的小距离（如10单位）
                            target_page.trimbox[2] - watermark_page.trimbox[2] ,
                            # 右上角水印左下角Y坐标，通常设为距离页面顶部边缘的小距离（如10单位）
                            target_page.trimbox[3] - watermark_page.trimbox[3] ,
                            # 右上角水印右上角X坐标，等于左下角X坐标加上水印宽度
                            target_page.trimbox[2] ,
                            # 右上角水印右上角Y坐标，等于左下角Y坐标加上水印高度
                            target_page.trimbox[3] 
                        )
                    )
        #保存添加了水印的PDF文件，并在原文件名基础上添加"已添加水印"的后缀
        marked_file = pdf_path.replace('.pdf', '_marked.pdf')
        target_pdf.save(marked_file)
        return marked_file

    def _add_score_watermark_by_postions(self, pdf_path, postions ) -> None:

        
        #打开目标PDF文件
        target_pdf = Pdf.open(pdf_path, allow_overwriting_input=True)

        #打开水印PDF文件并提取第一页作为水印
        watermark_pdf = Pdf.open(self.watermark_score_pdf_path)
        watermark_page = watermark_pdf.pages[0]

        #遍历目标PDF的每一页，根据行列数批量添加水印
        for idx,target_page in enumerate(target_pdf.pages):
            flag = False
            print('当前页：',idx)
            box = target_page.trimbox
            if idx in self.skip_pages:  # 如果当前页码在skip_pages列表中，则跳过
                continue
            for pos in postions:
                if pos['page'] == idx+1:
                    flag = True
                    #计算水印在目标页面上的位置并添加
                    #计算水印在目标页面上的位置并添加
                    index = pos['pos'].strip().split(' ')
                    index = [p for p in index if p!='']
                    x1,y1,x2,y2 = index
                    x1 = target_page.trimbox[2]-35 
                    # y1 = int(y1) - 10
                    x2 = target_page.trimbox[2]-10
                    # y2 = int(y2) + 10
                    pos['pos'] = f'{x1} {y1} {x2} {y2}'
            if flag:
                self.generator_score_postions_watermark(postions, idx+1, box[2], box[3])
                
                #打开水印PDF文件并提取第一页作为水印
                watermark_pdf = Pdf.open(self.watermark_score_pdf_path)
                watermark_page = watermark_pdf.pages[0]
                
                target_page.add_overlay(
                    watermark_page,
                    Rectangle(0,0,box[2],box[3])
                )
                # os.remove(self.watermark_score_pdf_path)

        #保存添加了水印的PDF文件，并在原文件名基础上添加"已添加水印"的后缀
        target_pdf.save(pdf_path.replace('.pdf','_with_watermark.pdf'))


    def make_score_watermark(self,pdf_path, teacher,score ,max_score):
        self.generator_score_watermark(teacher,score, max_score)
        marked_file = self._add_score_watermark(pdf_path)
        os.remove(self.watermark_pdf_path)
        return marked_file

    def make_score_watermark_by_postions(self, pdf_path, teacher, score, postions ):
        self.generator_score_watermark(teacher,score)
        self._add_score_watermark_by_postions(pdf_path,  postions)
        if os.path.exists(self.watermark_pdf_path):
            os.remove(self.watermark_pdf_path)
        

if __name__ == '__main__':
    pdf_path = '/root/k8s-vue/exam-platform-of-big-data/工具函数/autoExam_test.pdf'

    water = Watermark()
    water.make_score_watermark(pdf_path,'张三',90)
    
    
    pos = [{'page': 1, 'pos': '114 449 387 465 ', 'text': '第一题.(30分)编写程序，实现下列功能。', 'type': 'title', 'score': '', 'max_score': '30'}, {'page': 1, 'pos': '114 157 346 168 ', 'text': '"""代码_1.1：输出每位学员已通过的考试科目"""', 'type': 'code', 'score': '', 'max_score': ''}, {'page': 1, 'pos': '114 128 325 139 ', 'text': '"""代码_1.2：输出已通过驾考的学员名单"""', 'type': 'code', 'score': '', 'max_score': ''}, {'page': 1, 'pos': '114 74  509 104 ', 'text': '第二题.(40分)根据美国人口普查局1985年的工资普查数据集\nCPS1985.csv,编写程序绘制如下分析图：', 'type': 'title', 'score': '', 'max_score': '40'}, {'page': 3, 'pos': '114 428 378 439 ', 'text': '"""代码_2.1：绘制不同年龄员工平均工资变化折线图"""', 'type': 'code', 'score': '', 'max_score': ''}, {'page': 3, 'pos': '114 399 409 410 ', 'text': '"""代码_2.2：绘制不同年龄员工最低和最高工资对照柱状图"""', 'type': 'code', 'score': '', 'max_score': ''}, {'page': 3, 'pos': '114 370 388 381 ', 'text': '"""代码_2.3：绘制此次普查不同性别人数分布占比饼图"""', 'type': 'code', 'score': '', 'max_score': ''}, {'page': 3, 'pos': '114 316 512 346 ', 'text': '第三题.(30分) 按照指定网络结构建立卷积神经网络对下列猴\n痘病毒(monkeypox)图片进行猴痘与非猴痘分类。', 'type': 'title', 'score': '', 'max_score': '30'}, {'page': 5, 'pos': '114 239 294 250 ', 'text': '"""代码_3.1：划分训练集和测试集"""', 'type': 'code', 'score': '', 'max_score': ''}, {'page': 5, 'pos': '114 210 241 221 ', 'text': '"""代码_3.2：构建模型"""', 'type': 'code', 'score': '', 'max_score': ''}, {'page': 5, 'pos': '114 181 283 192 ', 'text': '"""代码_3.3：训练所构建的模型"""', 'type': 'code', 'score': '', 'max_score': ''}]
    water.make_score_watermark_by_postions(pdf_path,'张三',90)